import numpy as np
import matplotlib.pyplot as plt

def traj(t1, t2, x1, x2):
    A = np.array([
        [1, t1, t1**2, t1**3, t1**4, t1**5, t1**6, t1**7, t1**8, t1**9],
        [1, t2, t2**2, t2**3, t2**4, t2**5, t2**6, t2**7, t2**8, t2**9],
        [0, 1, 2*t1, 3*t1**2, 4*t1**3, 5*t1**4, 6*t1**5, 7*t1**6, 8*t1**7, 9*t1**8],
        [0, 1, 2*t2, 3*t2**2, 4*t2**3, 5*t2**4, 6*t2**5, 7*t2**6, 8*t2**7, 9*t2**8],
        [0, 0, 2, 6*t1, 12*t1**2, 20*t1**3, 30*t1**4, 42*t1**5, 56*t1**6, 72*t1**7],
        [0, 0, 2, 6*t2, 12*t2**2, 20*t2**3, 30*t2**4, 42*t2**5, 56*t2**6, 72*t2**7],
        [0, 0, 0, 6, 24*t1, 60*t1**2, 120*t1**3, 210*t1**4, 336*t1**5, 504*t1**6],
        [0, 0, 0, 6, 24*t2, 60*t2**2, 120*t2**3, 210*t2**4, 336*t2**5, 504*t2**6],
        [0, 0, 0, 0, 24, 120*t1, 360*t1**2, 840*t1**3, 1680*t1**4, 3024*t1**5],
        [0, 0, 0, 0, 24, 120*t2, 360*t2**2, 840*t2**3, 1680*t2**4, 3024*t2**5]
    ])

    B = np.array([x1, x2, 0, 0, 0, 0, 0, 0, 0, 0])

    # 求解线性方程组 A * C = B
    C = np.linalg.solve(A, B)
    return C


def man_desired(traj_theta_1, traj_theta_2, t):
    """
    计算给定时间 t 下的期望角度、角速度和角加速度。

    参数:
    traj_theta_1 (array-like): 关节1的9次多项式系数 (长度10)
    traj_theta_2 (array-like): 关节2的9次多项式系数 (长度10)
    t (float): 当前时间

    返回:
    theta_m (np.ndarray): [theta1, theta2] 的期望角度
    theta_dot_m (np.ndarray): [theta1_dot, theta2_dot] 的期望角速度
    theta_ddot_m (np.ndarray): [theta1_ddot, theta2_ddot] 的期望角加速度
    """

    # 构造时间向量
    T = np.array([t ** i for i in range(10)])
    # print(T)
    T_dot = np.array([0]+ [i * t ** (i - 1) for i in range(1, 10)])
    # print('T_dot:',T_dot)
    T_double_dot = np.array([0, 0] + [i * (i - 1) * t ** (i - 2) for i in range(2, 10)])
    # print(T_double_dot)

    # 计算期望角度、角速度和角加速度
    theta_m = np.array([
        np.dot(T, traj_theta_1),
        np.dot(T, traj_theta_2)
    ])

    theta_dot_m = np.array([
        np.dot(T_dot, traj_theta_1),
        np.dot(T_dot, traj_theta_2)
    ])

    theta_ddot_m = np.array([
        np.dot(T_double_dot, traj_theta_1),
        np.dot(T_double_dot, traj_theta_2)
    ])

    return theta_m, theta_dot_m, theta_ddot_m

if __name__ == '__main__':
    t1 = 0
    t2 = 10
    x1 = 0
    x2 = 20*np.pi/180
    C = traj(t1, t2, x1, x2)
    traj_theta_1 = traj(t1, t2, x1, 20*np.pi/180)
    traj_theta_2 = traj(t1, t2, x1, 20*np.pi/180)
    theta_ms, theta_dot_ms, theta_ddot_ms = [],[],[]
    # 生成时间序列
    t = np.linspace(t1, t2, 100)
    # 计算轨迹
    traj_values = sum(C[i] * t ** i for i in range(10))

    # 绘制轨迹
    plt.plot(t, traj_values, label='9 times traj')
    plt.scatter([t1, t2], [x1, x2], color='red', label='boundary')
    plt.xlabel('time')
    plt.ylabel('position')
    plt.legend()
    # plt.title('9次多项式轨迹规划')
    plt.grid(True)


    for i in t:
        print(i)
        theta_m, theta_dot_m, theta_ddot_m = man_desired(traj_theta_1,traj_theta_2,i)
        theta_ms.append(theta_m)
        theta_dot_ms.append(theta_dot_m)
        theta_ddot_ms.append(theta_ddot_m)
    plt.figure()
    print(theta_ms)
    plt.plot(t, theta_ms, label='9 times traj')
    plt.show()